# -*- coding: utf-8 -*-
"""
Created on 2021

@author: Ximing
"""
import base64
import hashlib
import random
import math


# 定义获取p,q,g的函数
def get_PQG():
    q = getPrime()
    i = 6
    while True:
        p = q * i + 1
        if is_prime(p):
            break
        i += 1
    h = random.randint(3, p - 2)
    # h = 2
    g1 = exp_mode(h, (p - 1) // q, p)
    return p, q, g1, h


# 定义获取公钥y的函数
def get_Y(g, x, p):
    return exp_mode(g, x, p)


# 定义获取素数的函数
def getPrime():
    # tmp = random.randint(Constant.MIN_Q, Constant.MAX_Q)
    ptmp = random.randint(1, 999999)
    while not is_prime(ptmp):
        ptmp = random.randint(1, 999999)
    return ptmp


# 判断是否为素数
def is_prime(n):
    if n <= 1:
        return False
    for i in range(2, int(math.sqrt(n)) + 1):
        if n % i == 0:
            return False
    return True


# 快速整数次幂取模(base ^ exponent) mod n
def exp_mode(base, exponent, n):
    bin_array = bin(exponent)[2:][::-1]
    r = len(bin_array)
    base_array = []
    pre_base = base
    base_array.append(pre_base)
    for _ in range(r - 1):
        next_base = (pre_base * pre_base) % n
        base_array.append(next_base)
        pre_base = next_base
    a_w_b = __multi(base_array, bin_array)
    return a_w_b % n


# 快速模乘
def __multi(array, bin_array):
    result = 1
    for index in range(len(array)):
        a = array[index]
        if not int(bin_array[index]):
            continue
        result *= a
    return result


# 定义签名函数 返回签名
def signature(p, q, g, sk, HM, k):
    r = exp_mode(exp_mode(g, k, p), 1, q)
    s = (inverseElement(k, q) * (HM + sk * r)) % q
    return r, s


# 求逆元
def inverseElement(a, b):
    # 将初始b的绝对值进行保存
    if b < 0:
        m = abs(b)
    else:
        m = b
    flag = gcd(a, b)

    # 判断最大公约数是否为1，若不是则没有逆元
    if flag == 1:
        r, x, y = ext_gcd(a, b)
        x0 = x % m  # 对于Python '%'就是求模运算，因此不需要'+m'
        return x0
    else:
        print(flag)
        print("Do not have!")


# 最大公约数
def gcd(a, b):
    if b == 0:
        return a
    else:
        return gcd(b, a % b)


# 扩展欧几里的算法：计算 ax + by = 1中的x与y的整数解（a与b互质），参数 a<b
def ext_gcd(a, b):
    if b == 0:
        x1 = 1
        y1 = 0
        x = x1
        y = y1
        r = a
        return r, x, y
    else:
        r, x1, y1 = ext_gcd(b, a % b)
        x = y1
        y = x1 - a // b * y1
        return r, x, y


# 定义DSA签名的验证函数
def verification(p, q, g, y, HM, r, s):
    print('开始验证签名')
    w = inverseElement(s, q)
    print(w)
    u1 = (HM * w) % q
    u2 = (r * w) % q
    vm1 = exp_mode(g, u1, p)  # (g ** u1) % p
    vm2 = exp_mode(y, u2, p)  # (y ** u2) % p
    v = ((vm1 * vm2) % p) % q
    return v


def SHA(filestr):
    print(filestr)
    # print("文件：" + filestr)
    s = hashlib.sha1()
    s.update(filestr)
    tmp = s.hexdigest()
    print("哈希值：" + tmp)
    return int(tmp, 16)


if __name__ == '__main__':
    p, q, g, h = get_PQG()
    print('全局参数 获取成功！')
    print('p = ' + str(p))
    print('q = ' + str(q))
    print('h = ' + str(h))
    print('g = ' + str(g))
    sk = random.randint(3, q)
    print('私钥生成成功，私钥是' + str(sk))
    pk = get_Y(g, sk, p)
    print('公开钥 pk = ' + str(pk))

    # 生成签名
    root = "dsa.py"
    print('要签名的文件是: ' + root)
    with open(root, "rb") as srcFile:
        tmp = base64.b64encode(srcFile.read())
    HM = SHA(tmp)
    # input('请输入为本次签名选取的秘密数 k , 其中 0 < k < ' + str(q) + '    ')
    k = random.randint(3, q)
    print('本次签名选取的秘密数是： ' + str(k))
    r, s = signature(p, q, g, sk, HM, k)
    print("签名成功，本条消息的签名(r,s)为 ( " + str(r) + '  ' + str(s) + ' )')

    # 验证签名
    with open(root, "rb") as srcFile:
        tmp = base64.b64encode(srcFile.read())
    HM = SHA(tmp)
    v = verification(p, q, g, pk, HM, r, s)
    print("v=", end="")
    print(v)
    if v == r:
        print("签名校验结果: Signature verification is true")
    else:
        print("签名校验结果: Sorry，signature verification is false")
